//
//  PoseOptimiser.cpp
//  markerARLib
//
//  Created by CADCG on 16/4/28.
//  Copyright © 2016年 CvidKal. All rights reserved.
//

#include "PoseOptimiser.hpp"
#include "Robost_Function.hpp"
#include <cmath>
#include <sophus/se3.h>
#include "ceres/ceres.h"
#include "glog/logging.h"
#include <ceres/rotation.h>
#include "Frame.hpp"
#include "LA.hpp"
#include "Algorithm.hpp"

#include <sophus/se3.h>
using ceres::AutoDiffCostFunction;
using ceres::CostFunction;
using ceres::Problem;
using ceres::Solver;
using ceres::Solve;

using ceres::SizedCostFunction;
using namespace Sophus;

namespace markerAR
{
    

namespace PoseOptimiser {
    
    
    
    
    
    
    class HomographyResidual:public ceres::SizedCostFunction<3,8>
    {
    public:
        HomographyResidual(double u1,double v1,double u2,double v2,const double *oldParam):u1(u1),u2(u2),v1(v1),v2(v2){memcpy(this->oldParam, oldParam, sizeof(double)*8);};
        virtual ~HomographyResidual(){};
        virtual bool Evaluate(double const* const* parameters,double *residuals,double **jacobians)const{
            const double * H = parameters[0];
            double z =  H[6]*u1+H[7]*v1+1;
            double x = H[0]*u1+H[1]*v1+H[2];
            double y = H[3]*u1+H[4]*v2+H[5];
            residuals[0] = x/z-u2;
            residuals[1] = y/z-v2;
            residuals[2] = 0;
//            memcpy((double*)r, (double*)residuals, sizeof(double)*2);
//            
            for (int i=0; i<8; ++i) {
                residuals[2] += 1/2*(oldParam[i]-H[i])*(oldParam[i]-H[i]);
            }
            if (jacobians != NULL && jacobians[0] != NULL) {
                jacobians[0][0] = u1/z; jacobians[0][1] = v1/z; jacobians[0][2] = 1/z; jacobians[0][3] = jacobians[0][4] = jacobians[0][5] = 0; jacobians[0][6] = x*(-1/z/z)*u1; jacobians[0][7] = x*(-1/z/z)*v1;
                jacobians[0][8] = jacobians[0][9] = jacobians[0][10] = 0; jacobians[0][11] = u1/z; jacobians[0][12] = v1/z; jacobians[0][13] = 1/z; jacobians[0][14] = y*(-1/z/z)*u1; jacobians[0][15] = y*(-1/z/z)*v1;
//                jacobians[0][16] = H[0]; jacobians[0][17] = H[1]; jacobians[0][18] = H[2]; jacobians[0][19] = H[3]; jacobians[0][20] = H[4]; jacobians[0][21] = H[5]; jacobians[0][22] = H[6]; jacobians[0][23] = H[7];
                memset(&jacobians[0][16], 0, sizeof(double)*8);
            }
//            memcpy((double*)this->oldParam, H, sizeof(double)*8);
            return true;
        }
        void computeJacobian(double x,double y,double z,const double*H,double **jacobians)
        {
            
        }
        
    private:
        double u1,v1,u2,v2;
        double oldParam[8];
        double r[2];
    };

    
    class HomographyProjection
    {
    public:
        HomographyProjection(double u1,double v1,double u2,double v2,const double *oldParam):u1(u1),u2(u2),v1(v1),v2(v2){memcpy(this->oldParam, oldParam, sizeof(double)*8);};
        virtual ~HomographyProjection(){};
        template <typename T>
        bool operator()(const T* const parameters,
                        T* residuals) const {
            const T * H = parameters;
            T z =  H[6]*T(u1)+H[7]*T(v1)+T(1);
            T x = H[0]*T(u1)+H[1]*T(v1)+H[2];
            T y = H[3]*T(u1)+H[4]*T(v2)+H[5];
            residuals[0] = x/z-T(u2);
            residuals[1] = y/z-T(v2);
            residuals[2] = T(0);

            return true;
        }
        
        
    private:
        double u1,v1,u2,v2;
        double oldParam[8];
    };
    
//    j =
//    [        (ry*x)/z^2, rz/z + (rx*x)/z^2, -ry/z, 1/z,   0, -x/z^2]
//    [ (ry*y)/z^2 - rz/z,        (rx*y)/z^2,  rx/z,   0, 1/z, -y/z^2]
    class BundlerResidual : public ceres::SizedCostFunction<3, 6> {
    public:
        BundlerResidual(double u,double v,double focal,double*pointXYZ,double*camera):u(u),v(v),focal(focal),pointXYZ(pointXYZ){memcpy(oldCamera,camera,sizeof(double)*6);}
        virtual ~BundlerResidual() {}
        virtual bool Evaluate(double const* const* parameters,
                              double* residuals,
                              double** jacobians) const {
            double p[3];
            const double * camera = parameters[0];
            ceres::AngleAxisRotatePoint(camera, pointXYZ, p);
            
            //apply translation
            p[0]+=camera[3];
            p[1]+=camera[4];
            p[2]+=camera[5];
            double x= p[0];
            double y= p[1];
            double z= p[2];
            
            double rx = x;
            double ry = y;
            double rz = z;

        
            double invz = 1/p[2];
            double invz2 = invz*invz;
        
            double xp =  focal * p[0] / p[2];
            double yp =  focal * p[1] / p[2];
            
            residuals[0] = xp-u;
            residuals[1] = yp-v;
            residuals[2] = 0;
            for (int i=0; i<6; i++) {
                residuals[2] += 1/2*(oldCamera[i]-camera[i])*(oldCamera[i]-camera[i]);
            }
            
            if (jacobians != NULL && jacobians[0] != NULL) {
                jacobians[0][0] = focal*ry*x*invz2; jacobians[0][1] =focal*(rz*invz2 + (rx*x)*invz2); jacobians[0][2]=-focal*ry/z; jacobians[0][3] =focal* invz; jacobians[0][4] = 0; jacobians[0][5] = focal*-x*invz2;
                jacobians[0][6] = focal*(ry*y*invz2-rz*invz); jacobians[0][7] = focal*rx*y*invz2; jacobians[0][8] = focal*rx*invz; jacobians[0][9]=0; jacobians[0][10] = focal*invz; jacobians[0][11] = focal*-y*invz2;
                jacobians[0][12] = camera[0];jacobians[0][13] = camera[1];jacobians[0][14] = camera[2];
                jacobians[0][15] = camera[3];jacobians[0][16] = camera[4];jacobians[0][17] = camera[5];
            }
            memcpy((double*)oldCamera,(const double*)camera,sizeof(double)*6);
            return  true;

        }
    private:
        double u,v;
        double focal;
        double * pointXYZ;
        double oldCamera[6];
    };
    
    struct Projection {
        // (u, v): the position of the observation with respect to the image
        // center point.
        Projection(double u,double v,double focal,double*pointXYZ): u(u),v(v),focal(focal),pointXYZ(pointXYZ) {}
        
        template <typename T>
        bool operator()(const T* const camera,
                        T* residuals) const {
            T p[3];
            T point[3];
            point[0] = T(pointXYZ[0]);
            point[1] = T(pointXYZ[1]);
            point[2] = T(pointXYZ[2]);
            
            ceres::AngleAxisRotatePoint(camera, point, p);
            
            // Add the translation vector
            p[0] += camera[3];
            p[1] += camera[4];
            p[2] += camera[5];
            
            // Compute the center of distortion.  The sign change comes from
            // the camera model that Noah Snavely's Bundler assumes, whereby
            // the camera coordinate system has a negative z axis.
            T xp =  focal * p[0] / p[2];
            T yp =  focal * p[1] / p[2];
            
            
            residuals[0] =  xp - T(u);
            residuals[1] =  yp - T(v);
            return true;
        }
        
        double u,v;
        double focal;
        double * pointXYZ;
    };
    


    
    void PoseOptimiseByCeres(Frame*pf,int idx,bool verbose)
    {
        ceres::Problem problem;
        vector<Feature*> correspondences;
        int num_correspondences = 0;
        for (int i=0; i<pf->fts.size(); ++i) {
            if (pf->fts[i]->pointXYZ!=nullptr) {
                ++num_correspondences;
                correspondences.push_back(pf->fts[i]);
            }
        }
        
        double* R = pf->C[idx].R.data();
        double camera[6];
        ceres::RotationMatrixToAngleAxis(R, camera);
        camera[3] = pf->C[idx].t[0];
        camera[4] = pf->C[idx].t[1];
        camera[5] = pf->C[idx].t[2];
        double fx = pf->K->fx;
        double cx = pf->K->cx;
        double cy = pf->K->cy;
        for (int i = 0; i < num_correspondences; ++i)
        {
            
            // Each Residual block takes a point and a camera as input and
            // outputs a 2 dimensional residual.
            ceres::CostFunction *cost_function = new BundlerResidual(correspondences[i]->_pos[0]-cx,correspondences[i]->_pos[1]-cy,fx,correspondences[i]->pointXYZ->pointXYZ.data(),camera);
//            ceres::CostFunction *cost_function = new AutoDiffCostFunction<Projection,2,6>(new Projection(correspondences[i]->_pos[0]-cx,correspondences[i]->_pos[1]-cy,fx,correspondences[i]->pointXYZ->pointXYZ.data()));
            
            // Each observation correponds to a pair of a camera and a point
            // which are identified by camera_index()[i] and
            // point_index()[i] respectively.
            problem.AddResidualBlock(cost_function, new ceres::HuberLoss(1.73), camera);
            
        }
        
        
        

        
        Solver::Options options_;
        
        options_.linear_solver_type = ceres::DENSE_QR;
        options_.max_num_iterations = 10;
        options_.function_tolerance = 1e-6;
        options_.gradient_tolerance = 1e-6;
        options_.parameter_tolerance = 1e-6;
        options_.minimizer_progress_to_stdout = false;
        Solver::Summary summary;
        
        Solve(options_, &problem, &summary);
       
        if (verbose) {
            std::cout << summary.BriefReport() << "\n";
        }
        ceres::AngleAxisToRotationMatrix(camera, R);
        pf->C[idx].t[0]=camera[3];
        pf->C[idx].t[1]=camera[4];
        pf->C[idx].t[2]=camera[5];
    }
    
    
    
    
//    
//    int PoseOptimiseLM(Frame* pf)
//    {
//        std::vector<Feature*> correspondences;
//        int nInitialPoints = 0;
//        for(int i=0;i<pf->fts.size();++i)
//        {
//            if (pf->fts[i]->pointXYZ !=nullptr) {
//                nInitialPoints++;
//                correspondences.push_back(pf->fts[i]);
//                pf->fts[i]->inlier = true;
//            }
//        }
//        double distThreshold = 3.;//outlier threshold
//        RobustKernelHuber huber(distThreshold);
//        std::vector<Eigen::Vector2d> residual(correspondences.size());
//        double initialRes =0.;
//        
//        Eigen::Matrix<double, 6, 1> B;
//        Eigen::Matrix<double, 1, 6> J;
//
//        B.setZero();
//        Eigen::Matrix<double, 6, 6> JtJ;
//        JtJ.setZero();
//        for(int i=0;i<correspondences.size();++i)
//        {
//            Eigen::Matrix<double, 3, 6> j1;
//            Eigen::Matrix<double,2, 3> j2;
//            Eigen::Matrix<double, 1, 6> j;
//            residual[i] = pf->K->project(pf->C.map(correspondences[i]->pointXYZ->pointXYZ, j1),j2) - correspondences[i]->_pos;
//            Eigen::Matrix<double, 1, 2> j3;
//            j3(0,0) = 2*residual[i].x();
//            j3(0,1) = 2*residual[i].y();
//            j = j3*j2*j1;
//            double e = residual[i].squaredNorm();
//            Eigen::Vector3f rho;
//            huber.robustify(e, rho);
//            initialRes += rho(0);
//            B += -rho[1]*j.transpose()*sqrt(e);
//            J += j;
//            JtJ += j.transpose()*j*rho(1);
//            
//        }
//        
//        cout<<"Initial res:"<<initialRes<<endl;
//        cout<<J<<endl;
//        Eigen::Matrix<double, 6, 1> delta;
//        double step = 0.0001;
//        double mu = 1e-2;
//        
//        SE3 oldPose(pf->C.R,pf->C.t);
//        Eigen::Matrix<double,6,6> I;
//        I.setIdentity();
//        JtJ += I*mu;
//        
//        SE3 newPose;
//        for (int i=0; i<10; ++i)
//        {
//
//            delta = JtJ.inverse()*B;
//            newPose = oldPose.exp(delta);
//            Pose newP(newPose.rotation_matrix(),newPose.translation());
//            double totalRes = 0;
//            B.setZero();
//            JtJ.setZero();
//            J.setZero();
//            for(int i=0;i<correspondences.size();++i)
//            {
//                Eigen::Matrix<double, 3, 6> j1;
//                Eigen::Matrix<double,2, 3> j2;
//                Eigen::Matrix<double, 1, 6> j;
//                residual[i] = pf->K->project(newP.map(correspondences[i]->pointXYZ->pointXYZ, j1),j2) - correspondences[i]->_pos;
//                Eigen::Matrix<double, 1, 2> j3;
//                j3(0,0) = 2*residual[i].x();
//                j3(0,1) = 2*residual[i].y();
//                j = j3*j2*j1;
//                double e = residual[i].squaredNorm();
//                Eigen::Vector3f rho;
//                huber.robustify(e, rho);
//                totalRes += rho(0);
//                B += -rho[1]*j.transpose()*sqrt(e);
//                JtJ += j.transpose()*j*rho(1);
//                J += j;
//            }
//            JtJ += I*mu;
//            cout<<J<<endl;
//            cout<<"new res:"<<totalRes<<endl;
//            oldPose = newPose;
//        }
//        
//        return 0;
//    }
    
    
    
    void HomographyOptimiseByCeres(const vector<Eigen::Vector2d>&src,const vector<Eigen::Vector2d>&dst,Eigen::Matrix3d &H,std::vector<int>&mask,bool verbose)
    {
        ceres::Problem problem;
        H /=H(2,2);
        Eigen::Matrix3d H_RowMajor = H.transpose();
        for (int i = 0; i < src.size(); ++i)
        {
            
            // Each Residual block takes a point and a camera as input and
            // outputs a 2 dimensional residual.
//            ceres::CostFunction *cost_function = new HomographyResidual(src[i].x(),src[i].y(),dst[i].x(),dst[i].y(),H_RowMajor.data());
                        ceres::CostFunction *cost_function = new AutoDiffCostFunction<HomographyProjection,3,8>(new HomographyProjection(src[i].x(),src[i].y(),dst[i].x(),dst[i].y(),H_RowMajor.data()));
            
            // Each observation correponds to a pair of a camera and a point
            // which are identified by camera_index()[i] and
            // point_index()[i] respectively.
            problem.AddResidualBlock(cost_function, new ceres::HuberLoss(3.0  ), H_RowMajor.data());
            
        }
        for(int i=0;i<src.size();++i)
        {
            double x2,y2;
            applyHomography(H_RowMajor.transpose(), src[i].x(), src[i].y(), x2, y2);
            double e = (x2-dst[i].x())*(x2-dst[i].x())+(y2-dst[i].y())*(y2-dst[i].y());
            if (e<9.0)
                mask.push_back(i);
            
        }
        
        
        Solver::Options options_;
        
        options_.linear_solver_type = ceres::DENSE_QR;
        options_.max_num_iterations = 10;
        options_.function_tolerance = 1e-6;
        options_.gradient_tolerance = 1e-6;
        options_.parameter_tolerance = 1e-6;
        options_.minimizer_progress_to_stdout = false;
        Solver::Summary summary;
        
        Solve(options_, &problem, &summary);
        
        
        
        if (verbose) {
            std::cout << summary.BriefReport() << "\n";
        }
        H = H_RowMajor.transpose();
    }
}
    
    
    
}